#!/usr/bin/python
# -*- coding:utf-8 -*-
# -------------------------------------------------------
# @author  : 龍皓
# @time    : 2025/3/11 17:23
# @function: 请求工具类。
# @version : 1.0.0
# -------------------------------------------------------

from core.util.file_util import FileUtil
import os
import random
from core.util.logger import setup_logger

logger = setup_logger()

class RequestUtil:

    def __init__(self, config_path=None):
        if config_path is None:
            config_path = os.path.join(os.path.dirname(__file__), '..', 'config', 'request_headers.json')

        self.config_path = config_path
        self.load_headers()

    def load_headers(self):
        """从JSON文件中加载请求头"""
        data = FileUtil.read_json_file(self.config_path)
        if data is not None:
            # 假设 JSON 文件中的数据是一个包含多个请求头配置的列表
            self.headers_list = data
            logger.info(f"Loaded headers from {self.config_path}")
        else:
            self.headers_list = []
            logger.warning(f"No headers configuration found at {self.config_path}. Using empty headers list.")

    def save_headers(self):
        """将请求头保存到JSON文件"""
        FileUtil.write_json_file(self.config_path, self.headers_list)

    def add_header(self, key, value):
        """添加或更新请求头"""
        for headers in self.headers_list:
            headers[key] = value
        logger.debug(f"Added header: {key}={value}")
        self.save_headers()

    def remove_header(self, key):
        """移除请求头"""
        for headers in self.headers_list:
            if key in headers:
                del headers[key]
        logger.debug(f"Removed header: {key}")
        self.save_headers()

    def set_headers(self, new_headers_list):
        """完全替换请求头"""
        self.headers_list = new_headers_list
        logger.info("Headers have been completely replaced.")
        self.save_headers()

    def get_headers(self):
        """随机获取一个请求头配置"""
        if self.headers_list:
            random_headers = random.choice(self.headers_list)
            logger.debug(f"Randomly selected headers: {random_headers}")
            return random_headers.copy()
        else:
            logger.warning("No headers available.")
            return {}

if __name__ == '__main__':
    headers_manager = RequestUtil()

    # 示例：随机获取请求头
    random_headers = headers_manager.get_headers()
    print(random_headers)
